# -*- coding: utf-8 -*-
import numpy as np
import train_model.settings as settings


def generate_random_poetry(tokenizer, model, s=''):
    """
    随机生成一首诗
    :param tokenizer: 分词器
    :param model: 用于生成古诗的模型
    :param s: 用于生成古诗的起始字符串，默认为空串
    :return: 一个字符串，表示一首古诗
    """
    # 将初始字符串转成token
    token_ids = tokenizer.encode(s)
    # 去掉结束标记[SEP]
    token_ids = token_ids[:-1]
    while len(token_ids) < settings.MAX_LEN:
        # 进行预测，只保留第一个样例（我们输入的样例数只有1）的、最后一个token的预测的、不包含[PAD][UNK][CLS]的概率分布
        output = model(np.array([token_ids, ], dtype=np.int32))
        _probas = output.numpy()[0, -1, 3:]
        del output
        # print(_probas)
        # 按照出现概率，对所有token倒序排列
        p_args = _probas.argsort()[::-1][:100]
        # 排列后的概率顺序
        p = _probas[p_args]
        # 先对概率归一
        p = p / sum(p)
        # 再按照预测出的概率，随机选择一个词作为预测结果
        target_index = np.random.choice(len(p), p=p)
        target = p_args[target_index] + 3
        # 保存
        token_ids.append(target)
        if target == 3:
            break
    return tokenizer.decode(token_ids)


def generate_acrostic(tokenizer, model, head):
    """
    随机生成一首藏头诗
    :param tokenizer: 分词器
    :param model: 用于生成古诗的模型
    :param head: 藏头诗的头
    :return: 一个字符串，表示一首古诗
    """
    # 使用空串初始化token_ids，加入[CLS]
    token_ids = tokenizer.encode('')
    token_ids = token_ids[:-1]
    # 标点符号，这里简单的只把逗号和句号作为标点
    punctuations = ['，', '。']
    punctuation_ids = {tokenizer.token_to_id(token) for token in punctuations}
    # 缓存生成的诗的list
    poetry = []
    # 对于藏头诗中的每一个字，都生成一个短句
    for ch in head:
        # 先记录下这个字
        poetry.append(ch)
        # 将藏头诗的字符转成token id
        token_id = tokenizer.token_to_id(ch)
        # 加入到列表中去
        token_ids.append(token_id)
        # 开始生成一个短句
        while True:
            # 进行预测，只保留第一个样例（我们输入的样例数只有1）的、最后一个token的预测的、不包含[PAD][UNK][CLS]的概率分布
            output = model(np.array([token_ids, ], dtype=np.int32))
            _probas = output.numpy()[0, -1, 3:]
            del output
            # 按照出现概率，对所有token倒序排列
            p_args = _probas.argsort()[::-1][:100]
            # 排列后的概率顺序
            p = _probas[p_args]
            # 先对概率归一
            p = p / sum(p)
            # 再按照预测出的概率，随机选择一个词作为预测结果
            target_index = np.random.choice(len(p), p=p)
            target = p_args[target_index] + 3
            # 保存
            token_ids.append(target)
            # 只有不是特殊字符时，才保存到poetry里面去
            if target > 3:
                poetry.append(tokenizer.id_to_token(target))
            if target in punctuation_ids:
                break
    return ''.join(poetry)


def generate_beam_search(tokenizer, model, s='', beam_width=3, max_len=settings.MAX_LEN):
    """
    使用 Beam Search 策略生成古诗
    :param tokenizer: 分词器
    :param model: 用于生成古诗的模型
    :param s: 用于生成古诗的起始字符串，默认为空串
    :param beam_width: Beam Search 的宽度
    :param max_len: 最大生成长度
    :return: 一个字符串，表示生成的古诗
    """
    # 将初始字符串转成token
    token_ids = tokenizer.encode(s)
    # 去掉结束标记[SEP]
    token_ids = token_ids[:-1]

    # 初始化候选序列列表，每个候选序列包含当前token序列和累积概率
    candidates = [(token_ids, 1.0)]  # (token序列, 累积概率)

    while len(candidates[0][0]) < max_len:
        new_candidates = []
        for seq, prob in candidates:
            # 预测下一个token的概率分布
            output = model(np.array([seq, ], dtype=np.int32))
            _probas = output.numpy()[0, -1, 3:]  # 去掉特殊标记[PAD][UNK][CLS]
            del output

            # 获取概率最高的beam_width个token
            top_indices = _probas.argsort()[-beam_width:][::-1]
            top_probs = _probas[top_indices]

            for idx, p in zip(top_indices, top_probs):
                new_seq = seq + [idx + 3]  # 将token索引加入序列
                new_prob = prob * p  # 更新累积概率
                new_candidates.append((new_seq, new_prob))

        # 按累积概率排序，保留概率最高的beam_width个候选序列
        candidates = sorted(new_candidates, key=lambda x: x[1], reverse=True)[:beam_width]

        # 如果某个候选序列以结束标记[SEP]结尾，则提前结束
        if any(seq[-1] == 3 for seq, _ in candidates):
            break

    # 选择概率最高的候选序列作为最终结果
    best_seq = candidates[0][0]
    return tokenizer.decode(best_seq)


def generate_top_k_sampling(tokenizer, model, s='', k=5, max_len=settings.MAX_LEN):
    """
    使用 Top-k Sampling 策略生成古诗
    :param tokenizer: 分词器
    :param model: 用于生成古诗的模型
    :param s: 用于生成古诗的起始字符串，默认为空串
    :param k: 采样的候选词数量
    :param max_len: 最大生成长度
    :return: 一个字符串，表示生成的古诗
    """
    # 将初始字符串转成token
    token_ids = tokenizer.encode(s)
    # 去掉结束标记[SEP]
    token_ids = token_ids[:-1]

    while len(token_ids) < max_len:
        # 预测下一个token的概率分布
        output = model(np.array([token_ids, ], dtype=np.int32))
        _probas = output.numpy()[0, -1, 3:]  # 去掉特殊标记[PAD][UNK][CLS]
        del output

        # 获取概率最高的k个token及其概率
        top_indices = _probas.argsort()[-k:][::-1]
        top_probs = _probas[top_indices]

        # 归一化概率
        top_probs = top_probs / np.sum(top_probs)

        # 随机选择下一个token
        next_token = np.random.choice(top_indices, p=top_probs) + 3
        token_ids.append(next_token)

        # 如果生成了结束标记[SEP]，则提前结束
        if next_token == 3:
            break

    return tokenizer.decode(token_ids)
